import { OllamaEmbeddings } from "@langchain/community/embeddings/ollama";
import { PDFLoader } from "langchain/document_loaders/fs/pdf";
import { Chroma } from "@langchain/community/vectorstores/chroma";
import { RecursiveCharacterTextSplitter } from "langchain/text_splitter";
import { PromptTemplate } from "@langchain/core/prompts";
import { ChatOllama } from "@langchain/community/chat_models/ollama";
import { createStuffDocumentsChain } from "langchain/chains/combine_documents";
import { StringOutputParser } from "@langchain/core/output_parsers";
import { ChromaClient } from 'chromadb'

const collectionName = "hcy-1"

// 实例化embedding模型对象
const embeddings = new OllamaEmbeddings({
    model: "nomic-embed-text", // default value
    baseUrl: "http://117.72.38.226:11434", // default value
    requestOptions: {
    useMMap: true,
    numThread: 6,
    // numGpu: 1,
    },
});

 // 实例化一个用于对话的大模型对象
const qwen = new ChatOllama({
    baseUrl: "http://117.72.38.226:11434", // Default value
    model: "qwen:4b", // Default value
});
 // 给大模型的提示信息，这个提示词的权重很大，
 // 例如，这里如果是让大模型总结，他就会先回答用户问题，然后开始总结，回答的内容很多。 
 // 需求是让模型中介还是
const prompt = PromptTemplate.fromTemplate(
    "请根据文档内容回答用户的问题: {context}"
);
 // 创建一个langchain调用链，将模型和提示词信息对象串联起来
const chain = await createStuffDocumentsChain({
    llm: qwen,
    outputParser: new StringOutputParser(),
    prompt,
});
// 提出问题，这个地方就是读取用户的输入
const question = "金币怎么获取?";
// 实例化一个数据库对象
const vectorStore = await Chroma.fromExistingCollection(
    embeddings, // embedding 模型，要和向量化编码时使用同一模型
    { collectionName: collectionName,
      url: "http://117.72.38.226:8000",
    } // 要查询的数据集，相当于数据库或者表命
);
// 先查询向量库，拿到背景故事,取最符合的5个片段
const storyChunks = await vectorStore.similaritySearch(question,5);
console.log(storyChunks)
console.log(storyChunks.length)
// const s3 = Date.now()
// // 让大模型作答，context中就是本轮回答的背景故事，chain中包含了 对话模型、提示词。
// const response=await chain.invoke({
//     context: storyChunks,
// });
// const s4 = Date.now()
// console.log(response)
// console.log("llm回答问题耗时：",s4-s3," 毫秒")